package com.cxyjy.cfyuiapi.service.impl;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.cxyjy.cfyuiapi.service.QueryHistoryService;
import com.cxyjy.cfyuiapi.service.ComfyUIDownloadService;
import com.cxyjy.cfyuiapi.service.VideoService;
import com.cxyjy.common.util.JObjectUtil;
import com.cxyjy.common.util.JRequestUtil;
import com.cxyjy.common.util.JResponse;
import com.cxyjy.common.util.ResultUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;

import javax.servlet.http.HttpServletResponse;
import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.net.URLConnection;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static com.cxyjy.common.constants.CfyuiConstants.*;
import static com.cxyjy.common.util.JObjectUtil.stringValue;
import static com.cxyjy.common.util.ResultCode.GET_WRONG_PROMPT_ID;
import static com.cxyjy.common.util.ResultCode.WORKFLOW_STILL_IN_PROGRESS;

/**
 * @Author: jony
 * @Date : 2024/11/20 14:13
 * @VERSION v1.0
 */
@Slf4j
@Service("ComfyUIDownloadService")
public class ComfyUIDownloadServiceImpl implements ComfyUIDownloadService {
    @Autowired
    QueryHistoryService queryHistoryService;
    @Autowired
    VideoService videoService;

    @Override
    public void downloadVideo(String promptId, HttpServletResponse httpServletResponse) throws IOException {
        URL url = new URL(stringValue(HISTORY_URL) + "/" + promptId);
        Object data = JRequestUtil.sendGetRequest(url, null);
        if (data == null) {
            return;
        }
        Map<String, Object> dataMap = JObjectUtil.toMap(data);
        Map<String, Map<String, List<String>>> outputsMap = new HashMap<>();
        try {
            Map<String, Object> historyMap = JObjectUtil.toMap(dataMap.get(promptId));
            outputsMap = (Map<String, Map<String, List<String>>>) historyMap.get("outputs");
        } catch (NullPointerException e) {
            return;
        }

        Path videoFilePath = Paths.get(outputsMap.get("108").get("text").get(0));//目前换脸完整视频从108号节点出
        String videoFileName = stringValue(videoFilePath.getFileName()).replace("\"", "");
        String outputVideoUrlStr = stringValue(DOWNLOAD_URL) + "?filename=" + videoFileName;
        URL outputVideoUrl = new URL(outputVideoUrlStr);
        URLConnection urlConnection = outputVideoUrl.openConnection();
        // 将输入流包装成InputStream
        try (InputStream inputStream = new BufferedInputStream(urlConnection.getInputStream())) {
            httpServletResponse.setContentType("video/mp4");
//            httpServletResponse.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + videoFileName + "\"");
            int contentLength = urlConnection.getContentLength();
            httpServletResponse.setContentLength(contentLength);
            httpServletResponse.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment;filename=\""+"FaceSwapVideo-"+promptId+".mp4\"");

            // 将文件内容写入响应输出流
            byte[] buffer = new byte[4096];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                httpServletResponse.getOutputStream().write(buffer, 0, bytesRead);
            }
            httpServletResponse.getOutputStream().flush();
        }
    }

    @Override
    public void downloadImage(String promptId, HttpServletResponse httpServletResponse) throws IOException{
        URL url = new URL(stringValue(HISTORY_URL) + "/" + promptId);
        Object data = JRequestUtil.sendGetRequest(url, null);
        if (data == null) {
            return;
        }
        JSONObject dataJO = JSON.parseObject(JSON.toJSONString(data));
        JSONObject outputsImagesJO = JSON.parseObject(JSON.toJSONString(dataJO.getJSONObject(promptId)
                .getJSONObject("outputs")
                .getJSONObject("125")
                .getJSONArray("images")
                .get(0)));
        String filename=outputsImagesJO.getString("filename");
        String type=outputsImagesJO.getString("type");
        String outputImageUrlStr = stringValue(DOWNLOAD_URL) + "?filename=" + filename+"&type="+type;
        URL outputImageUrl = new URL(outputImageUrlStr);
        URLConnection urlConnection = outputImageUrl.openConnection();
        urlConnection.setRequestProperty("Accept-Encoding","identity");// 强制禁用压缩
        // 将输入流包装成InputStream
        try (InputStream inputStream = new BufferedInputStream(urlConnection.getInputStream())) {
            httpServletResponse.setContentType("image/jpeg");
            int contentLength = urlConnection.getContentLength();
            log.info(stringValue(contentLength));
            httpServletResponse.setContentLength(contentLength);
            httpServletResponse.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment;filename=\""+"FaceSwapImage-"+promptId+".jpg\"");

            // 将文件内容写入响应输出流
            byte[] buffer = new byte[4096];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                httpServletResponse.getOutputStream().write(buffer, 0, bytesRead);
            }
            httpServletResponse.getOutputStream().flush();
        }
    }

    @Override
    public JResponse checkPromptId(String promptId) throws IOException {
        if (JObjectUtil.isEmpty(promptId) || promptId.equals("undefined")) {
            return ResultUtil.error(10002, "请先完成上传图片生成视频流程！");
        }
        //使用comfyui /history接口查询工作流的信息
        URL url = new URL(stringValue(HISTORY_URL) + "/" + promptId);
        Object data = JRequestUtil.sendGetRequest(url, null);
        /**
         * 当查询出来的结果为空时要判断是用户输入的promptId错误，还是视频仍在生成当中
         * queue_running是正在运行的工作流 queue_pending是在排队的工作流
         * 两者之中都没有才能判定是用户输入的promptId错误
         */
        if (JObjectUtil.isEmpty(data)) {
            Map<String, Object> queueMap = JObjectUtil.toMap(JRequestUtil.sendGetRequest(QUEUE_URL, null));
            List<List> runningList = JObjectUtil.jsonArrayToList(stringValue(queueMap.get("queue_running")), List.class);
            List<List> pendingList = JObjectUtil.jsonArrayToList(stringValue(queueMap.get("queue_pending")), List.class);
            if (!runningList.isEmpty()) {
                if (runningList.get(0).get(1).equals(promptId)) {
                    return ResultUtil.error(WORKFLOW_STILL_IN_PROGRESS);
                }
                if(!pendingList.isEmpty()){
                    for(List<String> promptList : pendingList){
                        if(promptList.get(1).equals(promptId)){
                            return ResultUtil.error(WORKFLOW_STILL_IN_PROGRESS);
                        }
                    }
                }
            }

            return ResultUtil.error(GET_WRONG_PROMPT_ID);
        }
        Map<String, Object> dataMap = JObjectUtil.toMap(data);
        Map<String, Map<String, List<String>>> outputsMap = new HashMap<>();
        try {
            Map<String, Object> historyMap = JObjectUtil.toMap(dataMap.get(promptId));
            outputsMap = (Map<String, Map<String, List<String>>>) historyMap.get("outputs");
        } catch (NullPointerException e) {
            return ResultUtil.error(WORKFLOW_STILL_IN_PROGRESS);
        }
        if (outputsMap.isEmpty()) return ResultUtil.error(WORKFLOW_STILL_IN_PROGRESS);
        return ResultUtil.justSuccess();
    }


    @Override
    public void downloadVideoTest(String promptId, HttpServletResponse httpServletResponse) throws IOException {
        URL url = new URL(stringValue(HISTORY_URL) + "/" + promptId);
        Object data = JRequestUtil.sendGetRequest(url, null);
        if (data == null) {
            return;
        }
        Map<String, Object> dataMap = JObjectUtil.toMap(data);
        Map<String, Map<String, Object>> outputsMap = new HashMap<>();
        try {
            Map<String, Object> historyMap = JObjectUtil.toMap(dataMap.get(promptId));
            outputsMap = (Map<String, Map<String, Object>>) historyMap.get("outputs");
        } catch (NullPointerException e) {
            return;
        }

        Map<String, Object> paramsMap = new HashMap<>();
        if (!outputsMap.isEmpty()) {
            //避免输出了多个视频造成的麻烦，只取第一个视频的信息
            Map.Entry<String, Map<String, Object>> entry = outputsMap.entrySet().iterator().next();
            Map<String, List> digitMap = (Map) entry.getValue();
            paramsMap = JObjectUtil.toMap(digitMap.get("gifs").get(0));
        }
        //        Object result = JRequestUtil.sendGetRequest(VIEW_URL, paramsMap);
        //用原始文件名，免得搞不清楚
        String fileName = stringValue(paramsMap.get("filename"));
        URL faceSwapVideoUrl = JRequestUtil.createGetRequestURL(VIEW_URL, paramsMap);
        URLConnection urlConnection = faceSwapVideoUrl.openConnection();
        // 将输入流包装成InputStream
        try (InputStream inputStream = new BufferedInputStream(urlConnection.getInputStream())) {
            httpServletResponse.setContentType(MediaType.APPLICATION_OCTET_STREAM_VALUE);
            httpServletResponse.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + fileName + "\"");
            // 将文件内容写入响应输出流
            byte[] buffer = new byte[4096];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                httpServletResponse.getOutputStream().write(buffer, 0, bytesRead);
            }
            httpServletResponse.getOutputStream().flush();
        }
    }
}
